import glob
import os

import numpy as np
import torch
from PIL import Image
from torch.utils.data import Dataset

from utils.dataset import patient
from utils.train import imageUtils


class Dataset(Dataset):
    def __init__(self,
                 data_dir,
                 image_size=(256, 256, 1),
                 mask_type='both',
                 alpha=500,
                 sigma=20,
                 fill_mode='nearest',
                 random_data=True,
                 img_norm_mean=None,
                 img_norm_std=None,
                 seed=None,
                 ):
        assert mask_type in ['inner', 'outer', 'both']
        self.image_size = image_size
        self.data_dir = data_dir
        self.mask_type = mask_type
        self.alpha = alpha
        self.sigma = sigma
        self.fill_mode = fill_mode
        self.random_data = random_data
        if random_data:
            assert img_norm_mean, 'please input img_norm_mean'
            assert img_norm_std, 'please input img_norm_std'
            self.img_norm_mean = img_norm_mean
            self.img_norm_std = img_norm_std
        self.seed = seed
        self.images, self.masks, self.file_names = self.load_images()
        self.len = len(self.images)
        self.initImages()

    def __getitem__(self, index):
        image = self.images[index]
        mask = self.masks[index]

        # stack image + mask together to simultaneously augment
        stacked = np.concatenate((image, mask), axis=2)

        if self.alpha != 0 and self.sigma != 0:
            stacked = imageUtils.random_elastic_deformation(
                stacked, self.alpha, self.sigma, self.fill_mode)

        augmented_image, augmented_mask = self.transform_both(stacked)

        augmented_image = np.array(augmented_image)
        augmented_mask = [np.array(augmented_mask)]
        seg_labels = imageUtils.to_categorical(augmented_mask).squeeze()

        return augmented_image, augmented_mask, seg_labels

    def __len__(self):
        return self.len

    def load_images(self):
        """Load all patient images and contours from TrainingSet, Test1Set or
        Test2Set directory. The directories and images are read in sorted order.

        Arguments:
          data_dir - path to data directory (TrainingSet, Test1Set or Test2Set)

        Output:
          tuples of (images, masks), both of which are 4-d tensors of shape
          (batchsize, height, width, channels). Images is uint16 and masks are
          uint8 with values 0 or 1.
        """
        patient_dirs = []
        for data_dir in self.data_dir:
            glob_search = os.path.join(data_dir, "patient*")
            patient_dirs += sorted(glob.glob(glob_search))
        if len(patient_dirs) == 0:
            raise Exception("No patient directors found in {}".format(self.data_dir))

        # load all images into memory (dataset is small)
        images = []
        inner_masks = []
        outer_masks = []
        file_names = []
        for patient_dir in patient_dirs:
            p = patient.PatientData(patient_dir)
            images += p.images
            inner_masks += p.endocardium_masks
            outer_masks += p.epicardium_masks
            file_names += p.file_names

        # reshape to account for channel dimension
        images = np.asarray(images)[:, :, :, None]
        if self.mask_type == 'inner':
            masks = np.asarray(inner_masks)[:, :, :, None]
        elif self.mask_type == 'outer':
            masks = np.asarray(outer_masks)[:, :, :, None]
        elif self.mask_type == 'both':
            # mask = 2 for endocardium, 1 for cardiac wall, 0 elsewhere
            masks = np.asarray(inner_masks) + np.asarray(outer_masks)
            masks = masks[:, :, :, None]

        return images, masks, file_names

    def initImages(self):
        # before: type(masks) = uint8 and type(images) = uint16
        # convert images to double-precision
        self.images = self.images.astype('float64')

        if self.seed is not None:
            np.random.seed(self.seed)

    def transform_both(self, stacked):
        img = stacked[:, :, 0]
        mask = np.round(stacked[:, :, 1])
        img = Image.fromarray(img.squeeze().astype(np.uint8)).convert('L')
        mask = Image.fromarray(mask.squeeze().astype(np.uint8)).convert('L')
        if self.random_data:
            transform = imageUtils.Compose([
                imageUtils.RandomHorizontalFlip(),
                imageUtils.RandomVerticalFlip(),
                imageUtils.Scale((self.image_size[0], self.image_size[1])),
                imageUtils.RandomRotation(15),
                imageUtils.ToTensor(),
                imageUtils.Normalize(self.img_norm_mean, self.img_norm_std),
            ])
        else:
            transform = imageUtils.Compose([
                imageUtils.ToTensor(),
                imageUtils.Normalize(self.img_norm_mean, self.img_norm_std),
            ])

        img, mask = transform(img, mask)

        return img, mask
